#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>

#include "sdataset.h"
#include "spdf.h"

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;

//initialize random number generator
MTRand_int32 irand(time(0));

const double PI=3.14159265359;


float DeltaR(float eta, float phi)
{
  if(fabs(phi)> PI)
    phi= PI-fabs(phi);
  return sqrt(eta*eta + phi*phi);
}

int main(int argc, char **argv)
{
  if(argc < 1+1)
    {
      cout<<"Usage: ./sherpa_PDF_output input.root"<<endl;
      throw;
    }

  cout<<"Outputs are automatically stored in PDF_*"<<endl;
  
  //open rootfile
  TChain chain("Ntuple");
  chain.Add(argv[1]);

  //declare variables
  const int maxsize=10;
  Int_t           size;
  Double_t         w;   //[size]
  Float_t         pt[maxsize];   //[size]
  Float_t         eta[maxsize];
  Float_t         phi[maxsize];
  Float_t         m[maxsize];   //[size]
  Float_t         t1[maxsize];
  Float_t         t2[maxsize];
  Float_t         t3[maxsize];


  //read the input file
  chain.SetBranchAddress("size", &size);
  chain.SetBranchAddress("w", &w);
  chain.SetBranchAddress("pt", pt);
  chain.SetBranchAddress("eta", eta);
  chain.SetBranchAddress("phi", phi);
  chain.SetBranchAddress("m", m);
  chain.SetBranchAddress("tau1", t1);
  chain.SetBranchAddress("tau2", t2);
  chain.SetBranchAddress("tau3", t3);


  DataSet DataSet_tau21_2jet(make_list<double>(0, 0),
			     make_list<double>(1.0,1.0),
			     make_list<double>(0.02,0.02));

  DataSet DataSet_tau21_2jet_dR(make_list<double>(0, 0, 1.0),
				make_list<double>(1.0,1.0,3*PI),
				make_list<double>(0.02,0.02,0.05));

  DataSet DataSet_tau21_2jet_mass(make_list<double>(0, 0, 0),
				  make_list<double>(1.0,1.0, 3.0),
				  make_list<double>(0.02,0.02,0.05));

  /*
  DataSet DataSet_mpt_2jet_mass(make_list<double>(0, 0, 0),
				make_list<double>(1.0,1.0, 3.0),
				make_list<double>(0.02,0.02,0.05));

  DataSet DataSet_mpt_2jet_dR(make_list<double>(0, 0, 0),
			      make_list<double>(1.0,1.0, 3*PI),
			      make_list<double>(0.02,0.02,0.05));

  DataSet DataSet_mpt_2jet(make_list<double>(0, 0),
			   make_list<double>(1.0,1.0),
			   make_list<double>(0.01,0.01));

  DataSet DataSet_mpt_2jet_pt1_dR(make_list<double>(0, 0, 0, 0),
				  make_list<double>(1.0,1.0, 3, 3*PI),
				  make_list<double>(0.025,0.025, 0.075, 0.15));

  DataSet DataSet_mpt_pt(make_list<double>(0, 0),
			 make_list<double>(1.0, 3.0),
			 make_list<double>(0.02, 0.05));


  DataSet DataSet_mpt12_pt12(make_list<double>(0, 0, 0, 0),
			     make_list<double>(1.0, 1.0, 3.0, 3.0),
			     make_list<double>(0.025, 0.025, 0.075, 0.075));
  
  */

  //changed to logarithmic scale
  DataSet DataSet_mpt_2jet_mass(make_list<double>(0, 0, 0),
				make_list<double>(3.0,3.0,3.0),
				make_list<double>(0.06,0.06,0.05));

  DataSet DataSet_mpt_2jet_dR(make_list<double>(0, 0, 0),
			      make_list<double>(3.0,3.0,3*PI),
			      make_list<double>(0.06,0.06,0.05));

  DataSet DataSet_mpt_2jet(make_list<double>(0, 0),
			   make_list<double>(3.0,3.0),
			   make_list<double>(0.03,0.03));

  DataSet DataSet_mpt_2jet_pt1_dR(make_list<double>(0, 0, 0, 0),
				  make_list<double>(3.0,3.0, 3, 3*PI),
				  make_list<double>(0.075,0.075, 0.075, 0.15));

  DataSet DataSet_mpt_pt(make_list<double>(0, 0),
			 make_list<double>(3.0, 3.0),
			 make_list<double>(0.06, 0.05));


  DataSet DataSet_mpt12_pt12(make_list<double>(0, 0, 0, 0),
			     make_list<double>(3.0, 3.0, 3.0, 3.0),
			     make_list<double>(0.075, 0.075, 0.075, 0.075));
  


  //do log pt
  DataSet DataSet_tau21_pt(make_list<double>(0, 0),
			   make_list<double>(1.0, 3.0),
			   make_list<double>(0.02, 0.05));

  //loop over events
  double total_weight=0;
  int nevt=0;
  for(Long64_t i=0; i<chain.GetEntries(); i++)
    {
      
      if(i%10000==0)
	cout<<"\r\033[K"<<(i*10000/chain.GetEntries())*0.01<<"% trained..."
	    <<flush;


      chain.GetEntry(i);
      if(size<2) continue;
      
      //some sanity cuts
      if(t1[0] < 1e-5) continue;
      if(t1[1] < 1e-5) continue;
      if(t2[0] < 1e-5) continue;
      if(t2[1] < 1e-5) continue;
      if(pt[0] < 50.) continue;
      if(pt[1] < 50.) continue;
      if(m[0]<20) continue;
      if(m[1]<20) continue;

      ++nevt;

      double dR=DeltaR(eta[0]-eta[1], phi[0]-phi[1]);
      
      TLorentzVector j1,j2;
      j1.SetPtEtaPhiM(pt[0],eta[0],phi[0],m[0]);
      j2.SetPtEtaPhiM(pt[1],eta[1],phi[1],m[1]);
      double mjj=(j1+j2).M();
      double logmjj=log(mjj/50.);

      double logpt1=log(pt[0]/50.);
      double logpt2=log(pt[1]/50.);
      /*
      double mpt1=m[0]/pt[0];
      double mpt2=m[1]/pt[1];
      */
      double mpt1=-log10(m[0]/pt[0]);
      double mpt2=-log10(m[1]/pt[1]);

      double t21_1=t2[0]/t1[0];
      double t21_2=t2[1]/t1[1];


      DataSet_tau21_pt.Fill(make_list<double>(t21_1, logpt1));
      DataSet_tau21_pt.Fill(make_list<double>(t21_2, logpt2));

      DataSet_mpt_pt.Fill(make_list<double>(mpt1,logpt1 ));
      DataSet_mpt_pt.Fill(make_list<double>(mpt2,logpt2 ));

      DataSet_mpt12_pt12.Fill
	(make_list<double>(mpt1, mpt2,
			   logpt1, logpt2));

      DataSet_mpt_2jet.Fill(make_list<double>(mpt1, mpt2));

      DataSet_tau21_2jet.Fill(make_list<double>(t21_1,t21_2));

      DataSet_tau21_2jet_dR.Fill(make_list<double>(t21_1,
						   t21_2,
						   dR) );

      DataSet_tau21_2jet_mass.Fill(make_list<double>(t21_1,
						     t21_2,
						     logmjj) );
      
      DataSet_mpt_2jet_mass.Fill(make_list<double>(mpt1,
						   mpt2,
						   logmjj) );
      
      DataSet_mpt_2jet_dR.Fill(make_list<double>(mpt1,
						 mpt2,
						 dR));


      DataSet_mpt_2jet_pt1_dR.Fill(make_list<double>(mpt1,
						     mpt2,
						     logpt1, dR));
      
      total_weight+=w;
    }

  cout<<endl;

  //here compute come PDF
  PDF PDF_tau21_2jet=DataSet_tau21_2jet.ComputePDF(true);
  PDF PDF_tau21_2jet_dR=DataSet_tau21_2jet_dR.ComputePDF(true);
  PDF PDF_tau21_2jet_mass=DataSet_tau21_2jet_mass.ComputePDF(true);
  PDF PDF_mpt_2jet=DataSet_mpt_2jet.ComputePDF(true);
  PDF PDF_mpt_2jet_mass=DataSet_mpt_2jet_mass.ComputePDF(true);

  PDF PDF_tau21_pt=DataSet_tau21_pt.ComputePDF(true);
  PDF PDF_mpt_pt=DataSet_mpt_pt.ComputePDF(true);
  PDF PDF_mpt12_pt12=DataSet_mpt12_pt12.ComputePDF(true);

  PDF PDF_mpt_2jet_dR=DataSet_mpt_2jet_dR.ComputePDF(true);
  PDF PDF_mpt_2jet_pt1_dR=DataSet_mpt_2jet_pt1_dR.ComputePDF(true);

  PDF_mpt_2jet.Write("PDF_mpt_2jet");
  PDF_mpt_2jet_pt1_dR.Write("PDF_mpt_2jet_pt1_dR");
  PDF_mpt_2jet_dR.Write("PDF_mpt_2jet_dR");
  PDF_mpt_2jet.Write("PDF_mpt_2jet");

  PDF_tau21_2jet_dR.Write("PDF_tau21_2jet_dR");
  PDF_tau21_2jet_mass.Write("PDF_tau21_2jet_mass");
  PDF_mpt_2jet_mass.Write("PDF_mpt_2jet_mass");
  PDF_mpt12_pt12.Write("PDF_mpt12_pt12");

  PDF_mpt_pt.Write("PDF_mpt_pt");
  PDF_tau21_pt.Write("PDF_tau21_pt");

  cout<<"Program ends."<<endl;
  return 1;
}
  
